import logging
import os
import cv2
import numpy as np

# 创建一个日志器
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)  # 设置日志级别为INFO


def cv_show(name, img):
    """显示图像"""
    if img is not None:
        cv2.imshow(name, img)
        cv2.waitKey(0)
    else:
        logger.warning(f"图像 {name} 无法显示，图像为空。")


def detectAndDescribe(image, scale_factor=0.5):
    """检测并描述图像特征点，先缩小图像再还原特征点位置"""
    if image is None:
        logger.error("输入图像为空，无法进行特征点检测！")
        return [], None, None

    # 1. 缩小图像
    small_image = cv2.resize(image, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)

    # 2. 将缩小后的图像转换成灰度图
    gray = cv2.cvtColor(small_image, cv2.COLOR_BGR2GRAY)

    # 3. 创建 SIFT 描述符对象
    descriptor = cv2.SIFT_create()

    # 4. 检测特征点并计算描述符
    kps, des = descriptor.detectAndCompute(gray, None)

    # 5. 如果有特征点，恢复它们在原始图像上的位置
    kps_float = np.float32([((kp.pt[0] / scale_factor), (kp.pt[1] / scale_factor)) for kp in kps]) if kps else []

    return kps, kps_float, des


def crop_black_area(image):
    """裁剪图像中黑色区域"""
    if image is None:
        logger.error("输入图像为空，无法进行裁剪！")
        return None

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
    x, y, w, h = cv2.boundingRect(thresh)
    return image[y:y + h, x:x + w] if w > 0 and h > 0 else image


def resize_and_rotate_image(image, scale_factor=1, rotate_code=cv2.ROTATE_90_COUNTERCLOCKWISE):
    """调整图像大小和旋转"""
    if image is None:
        logger.error("输入图像为空，无法进行大小调整和旋转！")
        return None

    resized_image = cv2.resize(image, None, fx=scale_factor, fy=scale_factor)
    rotated_image = cv2.rotate(resized_image, rotate_code)
    return rotated_image


def crop_image_to_size(image1, image2):
    """裁剪第二张图像，使其与第一张图像的大小一致"""
    if image1 is None or image2 is None:
        logger.error("图像为空，无法进行裁剪！")
        return None

    height1, width1 = image1.shape[:2]
    cropped_image2 = image2[:height1, :width1]
    return cropped_image2


def stitch_images(imageA, imageB, dir=1):
    """拼接两张图片，并裁剪去掉黑色区域"""
    if imageA is None or imageB is None:
        logger.error("输入图像为空，无法进行拼接！")
        return None

    # 根据 dir 判断旋转方向
    rotate = cv2.ROTATE_90_COUNTERCLOCKWISE if dir else cv2.ROTATE_90_CLOCKWISE

    imageA = resize_and_rotate_image(imageA, rotate_code=rotate)
    imageB = resize_and_rotate_image(imageB, rotate_code=rotate)

    if imageA is None or imageB is None:
        logger.error("调整后的图像为空，无法继续拼接！")
        return None

    cropped_imageB = crop_image_to_size(imageA, imageB)

    if cropped_imageB is None:
        logger.error("裁剪后的图像为空，无法继续拼接！")
        return None

    kpsA, kps_floatA, desA = detectAndDescribe(imageA)
    kpsB, kps_floatB, desB = detectAndDescribe(cropped_imageB)

    if desA is None or desB is None:
        logger.error("图像特征描述符为空，无法继续拼接！")
        return None

    # 创建一个暴力匹配器（BFMatcher），默认使用欧几里得距离度量
    matcher = cv2.BFMatcher()

    # 使用 knnMatch 方法进行最近邻匹配，返回两个最接近的匹配点
    # desB: 图像B的描述符，desA: 图像A的描述符，2 表示返回每个特征点的前两个匹配
    rawMatches = matcher.knnMatch(desB, desA, 2)

    # # 创建FLANN匹配器
    # index_params = dict(algorithm=1, trees=10)  # 使用KD树，设定树的数量
    # search_params = dict(checks=50)  # 设置检查次数，越大匹配越精确
    #
    # matcher = cv2.FlannBasedMatcher(index_params, search_params)
    #
    # # 使用 knnMatch 方法进行最近邻匹配
    # rawMatches = matcher.knnMatch(desB, desA, k=2)

    good = []
    matches = []
    for m in rawMatches:
        if len(m) == 2 and m[0].distance < 0.5 * m[1].distance:
            good.append(m)
            matches.append((m[0].trainIdx, m[0].queryIdx))

    # 检查匹配点的数量是否大于4，RANSAC方法要求至少有4个匹配点才能计算单应性矩阵
    if len(matches) > 4:
        # 从匹配点中提取图像A和图像B的关键点坐标
        ptsA = np.float32([kps_floatA[i] for (i, _) in matches])  # 从图像A的关键点中提取坐标
        ptsB = np.float32([kps_floatB[i] for (_, i) in matches])  # 从图像B的关键点中提取坐标

        # 使用RANSAC算法计算单应性矩阵H（将图像B映射到图像A的坐标系中）
        H, mask = cv2.findHomography(ptsB, ptsA, cv2.RANSAC, 10)

        # 使用单应性矩阵H对图像B进行透视变换，使其与图像A对齐
        result = cv2.warpPerspective(imageB, H, (imageB.shape[1] + imageA.shape[1], imageB.shape[0]))

        # 将图像A拷贝到result的左边部分，拼接两张图像
        result[0:imageA.shape[0], 0:imageA.shape[1]] = imageA

        # 调用自定义函数裁剪拼接图像的黑色区域（无数据区域）
        cropped_result = crop_black_area(result)

        # 如果裁剪后的图像为空，返回错误并停止
        if cropped_result is None:
            logger.error("裁剪后的拼接图像为空！")
            return None

        # 根据参数`dir`判断旋转方向
        if dir:
            # 顺时针旋转90度
            return cv2.rotate(cropped_result, cv2.ROTATE_90_CLOCKWISE)
        else:
            # 逆时针旋转90度
            return cv2.rotate(cropped_result, cv2.ROTATE_90_COUNTERCLOCKWISE)
    else:
        # 如果匹配点少于4个，无法进行拼接，返回错误
        logger.error("未找到4个以上匹配点，拼接图像失败！")
        return None


def stitch_images_in_folder(folder_path, dir=1):
    """将文件夹中的图像按文件名排序拼接成一张图像"""
    if not os.path.exists(folder_path):
        logger.error(f"文件夹 {folder_path} 不存在！")
        return None

    image_files = [f for f in os.listdir(folder_path) if f.endswith(('.jpg', '.png', '.jpeg'))]
    if not image_files:
        logger.error(f"文件夹 {folder_path} 中没有有效的图像文件！")
        return None

    if dir:
        image_files.sort(key=lambda x: int(os.path.splitext(x)[0]))  # 假设文件名是数字
    else:
        image_files.sort(key=lambda x: int(os.path.splitext(x)[0]), reverse=True)  # 假设文件名是数字

    imageA = cv2.imread(os.path.join(folder_path, image_files[0]))
    if imageA is None:
        logger.error(f"无法读取图像 {image_files[0]}！")
        return None

    imageAold = None
    for image_file in image_files[1:]:
        imageB = cv2.imread(os.path.join(folder_path, image_file))
        if imageB is None:
            logger.error(f"无法读取图像 {image_file}！")
            continue

        imageA = stitch_images(imageB, imageA, dir)
        # if dir:
        #     imageA = stitch_images(imageB, imageA, dir)
        # else:
        #     imageA = stitch_images(imageA, imageB, dir)

        if imageA is None:
            logger.error(f"拼接失败，跳过图像 {image_file}！")
            return imageAold
        else:
            imageAold = imageA

    return imageA


def stitch_images_in_folder_pro(folder_path):
    # 获取图像
    imageA = stitch_images_in_folder(folder_path, dir=1)
    imageB = stitch_images_in_folder(folder_path, dir=0)

    # 获取图像A下部2/3
    heightA = imageA.shape[0]
    imageA_cropped = imageA[int(1 / 3 * heightA):, :]

    # 获取图像B上部2/3
    heightB = imageB.shape[0]
    imageB_cropped = imageB[:int(2 / 3 * heightB), :]

    imageC = stitch_images(imageB_cropped, imageA_cropped)

    # 图像缩小
    imageC = cv2.resize(imageC, None, fx=1 / 3, fy=1 / 3, interpolation=cv2.INTER_LINEAR)

    return imageC


# 示例调用
if __name__ == "__main__":
    # imageA = cv2.imread(r'D:\Users\Administrator\Desktop\test_image\ST_image\001.jpg')
    # imageB = cv2.imread(r'D:\Users\Administrator\Desktop\test_image\ST_image\000.jpg')
    # imageC = stitch_images(imageA, imageB)
    #
    # if imageC is not None:
    #     imageC = cv2.resize(imageC, None, fx=1/3, fy=1/3)
    #     cv_show('Stitched Image', imageC)

    # folder_path = r"D:\Users\Administrator\Desktop\test_image\ST_image"
    # result = stitch_images_in_folder(folder_path, dir=1)
    # if result is not None:
    #     cv2.imwrite(r"..\stitched_result1.jpg", result)
    #     cv_show('Stitched Image from Folder', result)

    folder_path = r"D:\Users\Administrator\Desktop\VoltWatcherDemo\IS_Image"
    result = stitch_images_in_folder_pro(folder_path)
    cv2.imwrite(r"..\stitched_result.jpg", result)
    cv_show('Stitched Image from Folder', result)
